import numpy as np
from numpy import log10
import statsmodels.api as sm
import pandas as pd


def calculate_day_ols(day_df, index_symbols):
    # 要求必须是一天的数据，不再做天数的校验
    ssr_dict = {}
    direction_flag_dict = {}
    day_df.drop(day_df.loc[day_df["volume"] == 0].index, axis=0, inplace=True)
    for index_symbol in index_symbols:
        index_df = day_df.loc[day_df["index"] == index_symbol]
        if index_df.shape[0] < 4:
            ssr_dict[index_symbol] = 0
        else:
            # 先取成交量排前4的合约
            index_df.sort_values(by="volume", ascending=False, inplace=True)
            index_df = index_df.iloc[:4]
            # 4个合约按编号顺序（即到期顺序）排序
            index_df.sort_values(by="symbol", ascending=True, inplace=True)
            vt_symbols = index_df["vt_symbol"].tolist()
            date_delta_list = []
            first_contract_month = 0
            for vt_symbol in vt_symbols:
                # 每个合约要查到期日
                contract_id, _ = vt_symbol.split(".")
                contract_month = int(contract_id[-2:])
                if first_contract_month == 0:
                    month_gap = 0
                    first_contract_month = contract_month
                else:
                    month_gap = contract_month - first_contract_month
                    if month_gap < 0:
                        month_gap += 12
                date_delta_list.append(month_gap)
            direction_flag_dict[index_symbol] = True
            price_list = []
            first_price = 0
            last_price = 0
            last_direction = 0
            for price in index_df["close"].tolist():
                if first_price == 0:
                    first_price = price
                    last_price = price
                    price_list.append(0)
                else:
                    price_list.append(log10(price / first_price))
                    direction = 1 if price > last_price else -1
                    if price == last_price:
                        direction = last_direction
                    if direction != last_direction and last_direction != 0:
                        direction_flag_dict[index_symbol] = False
                    last_direction = direction
                    last_price = price
            y_price_series = np.asarray(price_list, dtype="float")
            x = sm.add_constant(np.asarray(date_delta_list, dtype="float"))
            result = sm.OLS(y_price_series, x).fit()
            ols = result.params[1]
            ols = pow(10, ols) - 1
            ssr_dict[index_symbol] = ols
    return ssr_dict, direction_flag_dict

def extract_sec_id(vt_symbol: str) -> str:
    """
    return sec_id
    """
    return vt_symbol[:2] if vt_symbol[1].isalpha() else vt_symbol[0]


def vt_symbol_to_index_symbol(vt_symbol):
    symbol_id, exchange_value = vt_symbol.split(".")
    sec_id = extract_sec_id(vt_symbol)
    index_id = f"{sec_id}99"
    return f"{index_id}.{exchange_value}"


if __name__ == '__main__':
    index_symbols = [
            "CF99.CZCE",
            "FG99.CZCE",
            "MA99.CZCE",
            "OI99.CZCE",
            "RM99.CZCE",
            "SF99.CZCE",
            "SM99.CZCE",
            "SR99.CZCE",
            "TA99.CZCE",
            "ZC99.CZCE",
            "a99.DCE",
            "c99.DCE",
            "cs99.DCE",
            "eb99.DCE",
            "eg99.DCE",
            "i99.DCE",
            "j99.DCE",
            "jm99.DCE",
            "l99.DCE",
            "m99.DCE",
            "pp99.DCE",
            "v99.DCE",
            "y99.DCE",
            "ag99.SHFE",
            "al99.SHFE",
            "au99.SHFE",
            "bu99.SHFE",
            "cu99.SHFE",
            "fu99.SHFE",
            "hc99.SHFE",
            "ni99.SHFE",
            "pb99.SHFE",
            "rb99.SHFE",
            "ru99.SHFE",
            "sn99.SHFE",
            "sp99.SHFE",
            "zn99.SHFE"
        ]
    data_df = pd.read_csv("D:\\roll return\\srr策略实盘跟踪\\20210315每日ols计算准确性\\daily_bar.csv", index_col="Unnamed: 0")
    data_df.reset_index(inplace=True)
    index_symbol_1 = []
    for vt_symbol in index_symbols:
        index_symbol_1.append(vt_symbol_to_index_symbol(vt_symbol))
    data_df = data_df.loc[data_df["index"].isin(index_symbol_1)]
    srr, direction = calculate_day_ols(data_df, index_symbols)
    print(srr)
